from utils.prompt import *
from utils.llm import *
from openai import OpenAI
import sys
import json
import os
from tqdm import tqdm

root_path = './data/'

#client = OpenAI(api_key=os.getenv('DEEPSEEKKEY'), base_url="https://api.deepseek.com")
#callllm = llm_handler(client,'deepseek-chat',dbmhandler=None,systemprompt = "Please analyze the provided historical analysis trajectory and extract the root cause analysis results in the specified format. ")
client = OpenAI(api_key=os.getenv('CSTCKEY'),base_url="https://uni-api.cstcloud.cn/v1")
callllm = llm_handler(client,dbmhandler=None,systemprompt = "Please analyze the provided historical analysis trajectory and extract the root cause analysis results in the specified format. ")

processed_uuids = []
try:
    with open(f'{root_path}/result/answer_regenerate.jsonl','r') as fp:
        processed_uuids = [json.loads(i)['uuid'] for i in fp.readlines()]
    print(f"{len(processed_uuids) } processed !")
except Exception as err:

    pass
if __name__ == '__main__':
    trace_file = sys.argv[1]
    with open(trace_file,'r') as fp:
        trace_data = fp.readlines()

    for tracestr in tqdm(trace_data):
        tracejson = json.loads(tracestr)
        uuid = tracejson['uuid']
        history = tracejson['trace']
        if uuid in processed_uuids: 
            continue
        input_str = gen_result(re.sub(r'(^#)|(\n#)|( +#)','\n##',history),uuid)
        written = False
        while not written:
            result = callllm(input_str).strip()
            print(result)
            try:
                result = json.loads(result[result.index('```')+7:-3])
                result['uuid'] = uuid 
                trace = json.dumps({'uuid':uuid,'trace':history})
                if 'uuid' in result.keys() and "component" in result.keys() and  "reason" in result.keys() and "reasoning_trace"  in result.keys():
                    with open(f'{root_path}/result/answer_regenerate.jsonl','a') as fp:
                        fp.write(json.dumps(result)+'\n')
                    written = True
            except Exception as err:
                print(err,'Rewritten')
